PYTHON-5903 Retry attempts should share a single stable operation_id by shrey-mongo · Pull Request #2901 · mongodb/mongo-python-driver

shrey-mongo · 2026-06-29T15:26:07Z

PYTHON-5903: Retry attempts should share a single stable operation_id

Changes in this PR

Command-monitoring events carry an operation_id to correlate the events of one logical operation. For single-document CRUD and cursor operations, each retry attempt got a different operation_id (it defaulted to the per-attempt request_id), so consumers couldn't group a retried operation's events. Bulk writes already shared one across batches.

_ClientConnectionRetryable mints the operation_id once and stamps it on the checked-out connection (conn.op_id) on every attempt
Connection gains an op_id attribute, reset on check-in.
pool.command and server.run_operation source op_id from the connection and pass it to command_runner.run_command / run_cursor_command, which thread it through to the publish_command_* events.

Test Plan

New test/asynchronous/test_operation_id_retry.py forces multiple retries via a failCommand failpoint across a read/write matrix and asserts all command events share one operation_id; verifies non-retryable multi-writes don't retry.
Ran the new tests plus the existing command-monitoring and retryable reads/writes suites against a single-node replica set which passed.

Checklist

Checklist for Author

Did you update the changelog (if necessary)?
Is there test coverage?
Is any followup work tracked in a JIRA ticket? If so, add link(s).

Checklist for Reviewer

Does the title of the PR reference a JIRA Ticket?
Do you fully understand the implementation? (Would you be comfortable explaining how this code works to someone else?)
Is all relevant documentation (README or docstring) updated?

NoahStapp · 2026-06-29T16:46:08Z

Please fill out the pull request template completely.

Copilot

Pull request overview

This PR ensures command-monitoring operation_id remains stable across retry attempts for retryable reads/writes and relevant cursor operations, improving APM correlation for a single logical operation.

Changes:

Mint a single per-operation operation_id in the retryable-operation wrapper and stamp it onto the checked-out connection for each attempt.
Thread the connection’s op_id through pool/server command execution into command_runner so published command-monitoring events use it.
Add sync + async integration tests that force multiple retries via failCommand and assert all related command events share one operation_id.

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
`tools/synchro.py`	Registers the new test file in the synchro conversion list.
`test/test_operation_id_retry.py`	Sync integration coverage asserting `operation_id` stability across retries (and no retry for non-retryable multi-writes).
`test/asynchronous/test_operation_id_retry.py`	Async equivalent of the retry `operation_id` stability tests.
`pymongo/synchronous/server.py`	Passes `conn.op_id` into cursor command execution for APM publishing.
`pymongo/synchronous/pool.py`	Adds `Connection.op_id`, forwards it into `run_command`, and clears it on check-in.
`pymongo/synchronous/mongo_client.py`	Mints/stabilizes per-operation ids across retries and stamps them on checked-out connections.
`pymongo/synchronous/command_runner.py`	Accepts/threads `op_id` through `run_command` / `run_cursor_command` into `_run_command`.
`pymongo/asynchronous/server.py`	Async counterpart: passes `conn.op_id` for cursor command execution.
`pymongo/asynchronous/pool.py`	Async counterpart: adds/forwards/clears `op_id` on connections.
`pymongo/asynchronous/mongo_client.py`	Async counterpart: stable per-operation id and per-attempt stamping onto connections.
`pymongo/asynchronous/command_runner.py`	Async counterpart: threads `op_id` through command execution to APM event publishing.

shrey-mongo · 2026-06-29T17:12:33Z

Updated description to include full template

codecov-commenter · 2026-06-29T18:21:09Z

Codecov Report

❌ Patch coverage is 97.72727% with 1 line in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
pymongo/_op_id.py	93.75%	0 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

NoahStapp · 2026-06-29T19:02:21Z

This needs to be updated to use the new operation_id, if present.

NoahStapp · 2026-06-30T15:05:07Z

There are merge conflicts, please resolve.

shrey-mongo · 2026-06-30T16:38:12Z

I fixed the merge conflict and the typing error that was causing the workflow to fail

NoahStapp · 2026-07-01T14:31:41Z

            func=func,
        )

-    def require_failCommand_appName(self, func):


What's the purpose of this change?

It was needed to fix an untyped-decorator mypy error that was occurring in my test. This matches the signature of require_failCommand_fail_point.

You don't need both the require_failCommand_fail_point and require_failCommand_appName annotations, the latter is a strict subset of the former.

NoahStapp · 2026-07-08T14:27:45Z

+
+
+class TestOperationIdRetry(AsyncIntegrationTest):
+    RETRIES = 5  # fail this many attempts; the (RETRIES + 1)th succeeds.


Why do we need 5 instead of 2 here? Does 5 provide additional test coverage or safety?

You're right, no added benefit, dropped to 2

NoahStapp · 2026-07-08T14:28:41Z

+        """Force ``retries`` retries of ``action`` and assert every command
+        event for ``command_name`` shares one integer operation_id."""
+        listener = AllowListEventListener(command_name)
+        client = await self.async_rs_or_single_client(event_listeners=[listener], appname=_APP_NAME)


Do we need to create a new client for every single operation, or can they all share a client set up in asyncSetUp?

Fixed, the client and listener are created once in asyncSetUp and shared across all operations

NoahStapp · 2026-07-08T14:30:07Z

+            with self.subTest(command=command_name):
+                listener = AllowListEventListener(command_name)
+                client = await self.async_rs_or_single_client(
+                    event_listeners=[listener], appname=_APP_NAME
+                )
+                coll = client.pymongo_test.test_operation_id_retry
+                await self._seed(coll)
+                listener.reset()
+
+                fail_point = {
+                    "mode": {"times": 1},
+                    "data": {
+                        "failCommands": [command_name],
+                        "closeConnection": True,
+                        "appName": _APP_NAME,
+                    },
+                }
+                async with self.fail_point(fail_point):
+                    with self.assertRaises(ConnectionFailure):
+                        await action(coll)


This looks very similar to what _check_stable_operation_id does, can we refactor them to reduce code duplication?

Refactored, both tests now call a shared _run_under_failpoint helper

NoahStapp · 2026-07-08T17:26:49Z

            async with operation.conn_mgr._lock:
                async with _MongoClientErrorHandler(self, server, operation.session) as err_handler:  # type: ignore[arg-type]
                    err_handler.contribute_socket(operation.conn_mgr.conn)
+                    operation.conn_mgr.conn.op_id = _randint()


This will assign a new op_id to every call here even though nothing retryable goes through this path.

Good point, I changed it to conn.op_id = None so each get More falls back to its own request_id.

NoahStapp · 2026-07-08T17:27:52Z

        db, coll = namespace.split(".", 1)
        spec = {"killCursors": coll, "cursors": cursor_ids}
+        # killCursors is its own op; drop any op_id left on a pinned cursor conn.
+        conn.op_id = None


A cleaner way of doing this might be to set conn.op_id = None in AsyncMongoClient._checkout to cover the edge cases and ensure that only the retryable logic sets a new op_id.

Done. added conn.op_id = None on checkout and removed the clear from checkin so every checkout starts at None and only the retryable read/write logic sets a real op_id.

NoahStapp · 2026-07-08T17:28:54Z

    more_to_come: bool = False,
    unpack_res: Optional[Callable[..., Any]] = None,
    cursor_id: Optional[int] = None,
+    op_id: Optional[int] = None,


Is this (and the one for run_command) ever used, since an op_id is set by the retryable read/write code in mongo_client.py?

You're right, removed the op_id param and they read from conn.op_id directly now.

NoahStapp · 2026-07-08T17:30:08Z

                user_fields=user_fields,
                exhaust_allowed=exhaust_allowed,
                write_concern=write_concern,
+                op_id=self.op_id,


We don't need to pass this since run_command is already passed the connection itself that owns op_id.

… into shrey-mongo/PYTHON-5903

NoahStapp · 2026-07-08T19:33:56Z

+                    prev_op_id = conn.op_id
+                    conn.op_id = None
                    await conn.authenticate(reauthenticate=True)
+                    conn.op_id = prev_op_id


This should use a try-except-finally block to ensure the restore of the previous op_id always occurs.

NoahStapp · 2026-07-08T19:41:38Z

                yield session._pinned_connection
                return
            async with await server.checkout(handler=err_handler) as conn:
+                conn.op_id = None  # Only retryable read/write logic sets an op_id.


This check should be moved to before the if in_txn block to prevent a pinned connection from re-using an old op_id.

Added reset on this branch

NoahStapp · 2026-07-08T19:42:49Z

+# Each operation, paired with the wire command it issues and an awaitable action.
+# These are all retryable; a stable operation_id must span every retry attempt.


This comment isn't needed: it describes a known standard pattern in the codebase.

NoahStapp

What do you think about using a ContextVar for the operationId instead of storing it on the connection? That would remove the current friction between a per-operation object being stored on a cross-operation object and simplify the cleanup significantly.

We already use this pattern in _csot.py and elsewhere for exactly the reason it might be useful here: to avoid threading through the entire call stack while still centralizing cleanup.

Something like this in _ClientConnectionRetryable.run:

token = _op_id.OP_ID.set(self._operation_id)
try:
   return await self._run()
finally:
   _op_id.OP_ID.reset(token)

And then replacing the current _handle_reauth block with this:

token = _op_id.OP_ID.set(None)
try:
   await conn.authenticate(reauthenticate=True)
finally:
   _op_id.OP_ID.reset(token)

The other cleanups could then be removed since they don't interact with _ClientConnectionRetryable.run() at all.

shrey-mongo · 2026-07-09T15:55:23Z

Implemented the ContextVar approach which helped with the cleanup, pool.py, cursor_base.py, and server.py match master and the explicit clears were removed for checkout, killCursors, and pinned getMore.

I set the token in _write/_read around self._func(...) instead of in run() since setting it in run would have the token live during connection checkout so the handshake/auth commands would inherit the operation's id.

… into shrey-mongo/PYTHON-5903 # Conflicts: # pymongo/asynchronous/mongo_client.py # pymongo/synchronous/mongo_client.py

… into shrey-mongo/PYTHON-5903

NoahStapp

OP_ID needs to be cleared inside AsyncPeriodicExecutor._run(), just as _csot.reset_all() is to prevent bleeding of op ids between async tasks.

Also look at _csot.py for an example of how to consolidate the OP_ID set/unset logic into a potential context manager.

NoahStapp · 2026-07-09T17:14:16Z

        codec_options=codec_options,
        user_fields=user_fields,
        orig=orig,
+        op_id=_op_id.OP_ID.get(),


We can consolidate these explicit passes of op_id to a new block within _run_command:

if op_id is None: op_id = _op_id.OP_ID.get()

NoahStapp · 2026-07-09T17:15:17Z

                if self._retrying:
                    self._log_retry(is_write=True)
-                return await self._func(self._session, conn, self._retryable)  # type: ignore
+                # Publish this op's id on every command of every attempt.


Suggested change

# Publish this op's id on every command of every attempt.

# One operation id across all attempts of this operation.

NoahStapp · 2026-07-09T17:15:34Z

            if self._retrying:
                self._log_retry(is_write=False)
-            return await self._func(self._session, self._server, conn, read_pref)  # type: ignore
+            # Publish this op's id on every command of every attempt.


Suggested change

# Publish this op's id on every command of every attempt.

# One operation id across all attempts of this operation.

NoahStapp · 2026-07-09T17:18:07Z

+
+
+class TestOperationIdRetry(AsyncIntegrationTest):
+    RETRIES = 2  # fail this many attempts; the (RETRIES + 1)th succeeds.


Is this comment adding something non-obvious to a reader of the code?

NoahStapp · 2026-07-09T17:21:45Z

+        """Force ``times`` closeConnection failures of ``command_name``, run
+        ``action``, and return its ``(started, failed, succeeded)`` events."""


Suggested change

"""Force ``times`` closeConnection failures of ``command_name``, run

``action``, and return its ``(started, failed, succeeded)`` events."""

"""Set a failpoint for the given command and return the corresponding events published during its execution."""

NoahStapp · 2026-07-09T17:26:50Z

+        await self.coll.insert_many([{"_id": i, "x": i % 3} for i in range(5)])
+        await self.coll.create_index("x")
+
+    async def _run_under_failpoint(self, command_name, action, times, expected_error=None):


Suggested change

async def _run_under_failpoint(self, command_name, action, times, expected_error=None):

async def _run_under_failpoint(self, name, f, times, expected_error=None):

NoahStapp · 2026-07-09T17:27:28Z

+    ("listIndexes", lambda c: _list_indexes(c)),
+]
+
+_COMMANDS = {name for name, _ in _RETRYABLE_WRITES + _RETRYABLE_READS}


If this isn't going to be reused, can it be defined inline at its use instead?

NoahStapp · 2026-07-09T17:28:36Z

+        for command_name, action in _RETRYABLE_WRITES:
+            with self.subTest(command=command_name):
+                await self._check_stable_operation_id(command_name, action, self.RETRIES)
+


Suggested change

for command_name, action in _RETRYABLE_WRITES:

with self.subTest(command=command_name):

await self._check_stable_operation_id(command_name, action, self.RETRIES)

for name, f in _RETRYABLE_WRITES:

with self.subTest(command=name):

await self._check_stable_operation_id(name, f, self.RETRIES)

We use the name, f convention for similar tests elsewhere.

NoahStapp · 2026-07-09T17:29:30Z

+        for command_name, action in _RETRYABLE_READS:
+            with self.subTest(command=command_name):
+                await self._check_stable_operation_id(command_name, action, self.RETRIES)


Suggested change

for command_name, action in _RETRYABLE_READS:

with self.subTest(command=command_name):

await self._check_stable_operation_id(command_name, action, self.RETRIES)

for name, f in _RETRYABLE_READS:

with self.subTest(command=name):

await self._check_stable_operation_id(name, f, self.RETRIES)

NoahStapp · 2026-07-09T18:22:18Z

+    async def test_non_retryable_write_is_not_retried(self):
+        # Multi-document writes are not retryable: a single network error must
+        # surface immediately, with exactly one attempt.
+        for command_name, action in [


Suggested change

for command_name, action in [

for name, f in [

… into shrey-mongo/PYTHON-5903

shrey-mongo requested a review from a team as a code owner June 29, 2026 15:26

shrey-mongo requested a review from aclark4life June 29, 2026 15:26

shrey-mongo force-pushed the shrey-mongo/PYTHON-5903 branch 2 times, most recently from 247b26f to 682dbac Compare June 29, 2026 15:56

NoahStapp requested review from NoahStapp and Copilot and removed request for aclark4life June 29, 2026 16:44

Copilot started reviewing on behalf of NoahStapp June 29, 2026 16:44 View session

Copilot AI reviewed Jun 29, 2026

View reviewed changes

NoahStapp requested changes Jun 29, 2026

View reviewed changes

shrey-mongo requested a review from NoahStapp June 30, 2026 15:04

fixed merge conflict

75f4e21

shrey-mongo force-pushed the shrey-mongo/PYTHON-5903 branch from 5546477 to 75f4e21 Compare June 30, 2026 15:46

NoahStapp requested changes Jul 1, 2026

View reviewed changes

fixed stale op id for reauth and killCursors ops and added justification

7237bf1

shrey-mongo requested a review from NoahStapp July 7, 2026 19:10

PYTHON-5903 merged master and resolved op_id conflicts

63312b0

shrey-mongo force-pushed the shrey-mongo/PYTHON-5903 branch from 1707aa0 to 63312b0 Compare July 8, 2026 17:23

NoahStapp requested changes Jul 8, 2026

View reviewed changes

shrey-mongo added 2 commits July 8, 2026 18:11

made requested changes

643ccef

Merge branch 'master' of https://github.com/mongodb/mongo-python-driver…

11223ce

… into shrey-mongo/PYTHON-5903

shrey-mongo requested a review from NoahStapp July 8, 2026 18:43

NoahStapp requested changes Jul 8, 2026

View reviewed changes

made requested changes

b36e20b

shrey-mongo requested a review from NoahStapp July 8, 2026 20:33

NoahStapp requested changes Jul 9, 2026

View reviewed changes

Comment thread test/asynchronous/test_operation_id_retry.py

shrey-mongo added 3 commits July 9, 2026 15:56

implemented ContextVar approach for op_id

d8f29b9

Merge branch 'master' of https://github.com/mongodb/mongo-python-driver…

371b161

… into shrey-mongo/PYTHON-5903 # Conflicts: # pymongo/asynchronous/mongo_client.py # pymongo/synchronous/mongo_client.py

Merge branch 'master' of https://github.com/mongodb/mongo-python-driver…

be5fa6e

… into shrey-mongo/PYTHON-5903

shrey-mongo requested a review from NoahStapp July 9, 2026 16:21

NoahStapp requested changes Jul 9, 2026

View reviewed changes

shrey-mongo added 2 commits July 9, 2026 19:01

made requested changes

dc74bf1

Merge branch 'master' of https://github.com/mongodb/mongo-python-driver…

f1e4285

… into shrey-mongo/PYTHON-5903

Jibola changed the title ~~PYTHON-5903 : Retry attempts should share a single stable operation_id~~ PYTHON-5903 Retry attempts should share a single stable operation_id Jul 9, 2026

shrey-mongo requested a review from NoahStapp July 9, 2026 19:18



		class TestOperationIdRetry(AsyncIntegrationTest):
		RETRIES = 5 # fail this many attempts; the (RETRIES + 1)th succeeds.

		# Each operation, paired with the wire command it issues and an awaitable action.
		# These are all retryable; a stable operation_id must span every retry attempt.

	# Publish this op's id on every command of every attempt.
	# One operation id across all attempts of this operation.



		class TestOperationIdRetry(AsyncIntegrationTest):
		RETRIES = 2 # fail this many attempts; the (RETRIES + 1)th succeeds.

		"""Force ``times`` closeConnection failures of ``command_name``, run
		``action``, and return its ``(started, failed, succeeded)`` events."""

	"""Force ``times`` closeConnection failures of ``command_name``, run
	``action``, and return its ``(started, failed, succeeded)`` events."""
	"""Set a failpoint for the given command and return the corresponding events published during its execution."""

	async def _run_under_failpoint(self, command_name, action, times, expected_error=None):
	async def _run_under_failpoint(self, name, f, times, expected_error=None):

Uh oh!

Conversation

shrey-mongo commented Jun 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Changes in this PR

Test Plan

Checklist

Checklist for Author

Checklist for Reviewer

Uh oh!

NoahStapp commented Jun 29, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

shrey-mongo commented Jun 29, 2026

Uh oh!

codecov-commenter commented Jun 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

NoahStapp commented Jun 30, 2026

Uh oh!

shrey-mongo commented Jun 30, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shrey-mongo commented Jun 29, 2026 •

edited

Loading

codecov-commenter commented Jun 29, 2026 •

edited

Loading